using System;
using System.Collections.Generic;
using System.IO;
using UnityEditor;
using UnityEngine;
using UnityEngine.Rendering;

// ReSharper disable once CheckNamespace
public class ShaderComplexityEditor : EditorWindow
{
    private string CompileDestPath
    {
        get
        {
            string projectPath = Application.dataPath.Replace("Assets", "");
            string destPath = projectPath + "/TempShaderCompiled";
            return destPath;
        }
    }
        
    private static ShaderComplexityEditor m_Instance = null;
    private string m_CompiledShaderPath;
    private string m_OCResult = string.Empty;
    private List<ShaderInfo> m_ShaderInfos = new List<ShaderInfo>();
    private string m_Architecture = "Mali-T720";
    
    private static string s_ProjectDirectory = string.Empty;
    private Vector2 m_ScrollPos = Vector2.zero;
    private int m_ShowMode = 0;
    private string m_SearchPattern;
    List<string> m_SearchPatterns = new List<string>();
    
    private static string[] m_ArchitectureList = new[]
    {
        // Arm 5th Generation architecture
        "Immortalis-G720", "Mali-G720", "Mali-G620",
        // Valhall architecture
        "Mali-G715", "Mali-G710", "Mali-G615", "Mali-G610", "Mali-G510", "Mali-G310", "Mali-G78AE", "Mali-G78", "Mali-G77", "Mali-G68", "Mali-G57",
        // Bifrost architecture
        "Mali-G76", "Mali-G72", "Mali-G71", "Mali-G52", "Mali-G51", "Mali-G31",
        // Midgard architecture
        "Mali-T880", "Mali-T860", "Mali-T830", "Mali-T820", "Mali-T760", "Mali-T720"
    };
    private enum ShaderSourceType
    {
        GLSL,
        Shader,
        GameObject,
        Material,
    }
    private ShaderSourceType m_ShaderSourceType = ShaderSourceType.Shader;
        
    private GameObject m_SelectedGameObject;
    private List<Material> m_MaterialsOnGameObject = new List<Material>();

    private Material m_SelectedMaterial;
    private Shader m_SelectedShader;
    private bool m_IncludeAllVariants;

    [MenuItem("Tools/资源合规/Shader复杂度计算器")]
    private static void Open()
    {
        if (m_Instance != null)
        {
            m_Instance.Close();
            m_Instance = null;
        }

        m_Instance = GetWindow<ShaderComplexityEditor>();
        m_Instance.titleContent = new GUIContent("Shader复杂度计算器");
        m_Instance.minSize = new Vector2(550.0f, 600.0f);
        m_Instance.Show();
    }

    private void OnEnable()
    {
        int assetIndex = Application.dataPath.LastIndexOf("Assets", StringComparison.Ordinal);
        s_ProjectDirectory = Application.dataPath.Substring(0, assetIndex).Replace('\\', '/');
    }

    private void OnDestroy()
    {
        DirectoryInfo destDir = new DirectoryInfo(CompileDestPath);
        if (destDir.Exists)
        {
            if (EditorUtility.DisplayDialog("提示", "是否清空所有临时文件？", "清空", "取消"))
            {
                destDir.Delete(true);
                if (File.Exists(CompileDestPath + ".meta"))
                {
                    File.Delete(CompileDestPath + ".meta");
                }
            }
        }
    }

    private void OnGUI()
    {
        EditorGUILayout.BeginVertical("TextArea");
        EditorGUILayout.Space();
        string[] shaderSrcTexts = new[]
        {
            "编译后的GLES 3.0 Shader",
            "指定Shader",
            "指定GameObject",
            "指定材质"
        };
        m_ShaderSourceType = (ShaderSourceType) GUILayout.Toolbar((int) m_ShaderSourceType, shaderSrcTexts);

        switch (m_ShaderSourceType)
        {
            case ShaderSourceType.GLSL:
            {
                EditorGUILayout.BeginHorizontal();
                EditorGUILayout.TextField("编译后的GLES 3.0 Shader", m_CompiledShaderPath);
                if (GUILayout.Button("打开", GUILayout.Width(44.0f)))
                {
                    string path = EditorUtility.OpenFilePanel("打开", s_ProjectDirectory, "txt");
                    if (!string.IsNullOrEmpty(path))
                    {
                        m_CompiledShaderPath = path;
                    }
                    m_SearchPattern = string.Empty;
                    m_SearchPatterns.Clear();
                }
                EditorGUILayout.EndHorizontal();
                break;
            }
            case ShaderSourceType.Shader:
            {
                Shader newSelectd = EditorGUILayout.ObjectField("指定Shader", m_SelectedShader, typeof(Shader), true) as Shader;
                if (newSelectd != m_SelectedShader)
                {
                    m_SearchPattern = string.Empty;
                    m_SearchPatterns.Clear();
                    m_SelectedShader = newSelectd;
                }
                break;
            }
            case ShaderSourceType.GameObject:
            {
                GameObject newSelectd = EditorGUILayout.ObjectField("指定GameObject", m_SelectedGameObject, typeof(GameObject), true) as GameObject;
                if (newSelectd != m_SelectedGameObject)
                {
                    m_MaterialsOnGameObject.Clear();
                    m_SelectedGameObject = newSelectd;
                    if (m_SelectedGameObject != null)
                    {
                        if (PrefabUtility.IsPartOfPrefabAsset(m_SelectedGameObject) &&
                            EditorUtility.DisplayDialog("提示", "检测到选中GameObject为Prefab资源，是否需要实例化出来？", "确定", "取消"))
                        {
                            m_SelectedGameObject = PrefabUtility.InstantiatePrefab(m_SelectedGameObject) as GameObject;
                        }

                        if (m_SelectedGameObject == null)
                        {
                            break;
                        }

                        Renderer r = m_SelectedGameObject.GetComponent<Renderer>();
                        if (r == null)
                        {
                            break;
                        }

                        m_MaterialsOnGameObject.AddRange(r.sharedMaterials);
                    }
                }

                EditorGUILayout.BeginHorizontal();
                GUILayout.Space(16.0f);
                float origWidth = EditorGUIUtility.labelWidth;
                EditorGUIUtility.labelWidth = 48.0f;
                EditorGUILayout.BeginVertical();
                for (int i = 0; i < m_MaterialsOnGameObject.Count; ++i)
                {
                    Material mat = m_MaterialsOnGameObject[i];
                    EditorGUILayout.BeginHorizontal();
                    EditorGUILayout.ObjectField($"材质[{i}]", mat, typeof(Material), true);
                    if (GUILayout.Button("选中", GUILayout.Width(44.0f)))
                    {
                        m_SelectedMaterial = mat;
                        m_ShaderSourceType = ShaderSourceType.Material;
                    }

                    EditorGUILayout.EndHorizontal();
                }

                EditorGUILayout.EndVertical();
                EditorGUIUtility.labelWidth = origWidth;
                EditorGUILayout.EndHorizontal();
                break;
            }
            case ShaderSourceType.Material:
            {
                Material newMat = EditorGUILayout.ObjectField("材质", m_SelectedMaterial, typeof(Material), true) as Material;
                m_IncludeAllVariants = EditorGUILayout.Toggle("是否包含所有变体", m_IncludeAllVariants);
                if (m_SelectedMaterial != newMat)
                {
                    m_SelectedMaterial = newMat;
                    
                    string keywords = string.Empty;
                    if (m_SelectedMaterial != null)
                    {
                        bool first = true;
                        foreach (LocalKeyword keyword in m_SelectedMaterial.enabledKeywords)
                        {
                            if (first)
                            {
                                first = false;
                            }
                            else
                            {
                                keywords += " ";
                            }
                            keywords += keyword.name;
                        }
                    }
                    m_SearchPattern = keywords;
                }

                break;
            }
        }
        EditorGUILayout.Space();
        EditorGUILayout.EndVertical();
        
        int selectIndex = 0;
        for (int i = 0; i < m_ArchitectureList.Length; ++i)
        {
            string arch = m_ArchitectureList[i];
            if (arch == m_Architecture)
            {
                selectIndex = i;
                break;
            }
        }

        selectIndex = EditorGUILayout.Popup("GPU架构", selectIndex, m_ArchitectureList);
        m_Architecture = m_ArchitectureList[selectIndex];

        if (m_ShaderSourceType == ShaderSourceType.GLSL && !string.IsNullOrEmpty(m_CompiledShaderPath) && GUILayout.Button("分析", GUILayout.Height(32.0f)))
        {
            m_OCResult = string.Empty;
            m_ShaderInfos.Clear();
            m_SearchPatterns.Clear();
            
            EditorUtility.DisplayProgressBar("提示", "复杂度分析中...", 0.0f);

            ShaderComplexityUtils.CompileGLSLRequest request = ShaderComplexityUtils.RequestCompileGLSL(m_CompiledShaderPath, CompileDestPath);
            request.onComplete = (shaderPath, tempPath, error) =>
            {
                EditorUtility.ClearProgressBar();
                
                if (!string.IsNullOrEmpty(error))
                {
                    return;
                }
                List<ShaderInfo> shaderInfos = ShaderComplexityUtils.CompileToShaderInfos(null, shaderPath, m_Architecture, out string result);
                m_ShaderInfos = shaderInfos;
                m_OCResult = result;
            };
        }
        if (m_ShaderSourceType == ShaderSourceType.Shader && m_SelectedShader != null && GUILayout.Button("分析", GUILayout.Height(32.0f)))
        {
            m_OCResult = string.Empty;
            m_ShaderInfos.Clear();
            // m_SearchPattern = string.Empty;
            m_SearchPatterns.Clear();
            
            EditorUtility.DisplayProgressBar("提示", "复杂度分析中...", 0.0f);
            
            ShaderComplexityUtils.CalculateShaderComplexity(m_SelectedShader, m_Architecture, CompileDestPath,
                (shaderInfos, result, success) =>
                {
                    if (success)
                    {
                        m_ShaderInfos = shaderInfos;
                        m_OCResult = result;
                        Repaint();
                    }
                    EditorUtility.ClearProgressBar();
                });
        }
        if (m_ShaderSourceType == ShaderSourceType.Material && m_SelectedMaterial != null && GUILayout.Button("分析", GUILayout.Height(32.0f)))
        {
            m_OCResult = string.Empty;
            m_ShaderInfos.Clear();
            m_SearchPatterns.Clear();
            
            EditorUtility.DisplayProgressBar("提示", "复杂度分析中...", 0.0f);

            if (m_IncludeAllVariants)
            {
                ShaderComplexityUtils.CalculateShaderComplexity(m_SelectedMaterial.shader, m_Architecture, CompileDestPath,
                    (shaderInfos, result, success) =>
                    {
                        if (success)
                        {
                            m_ShaderInfos = shaderInfos;
                            m_OCResult = result;
                            Repaint();
                        }
                        EditorUtility.ClearProgressBar();
                    });
            }
            else
            {
                ShaderComplexityUtils.CalculateMaterialComplexity(m_SelectedMaterial, m_Architecture, CompileDestPath,
                    (shaderInfos, result, success) =>
                    {
                        if (success)
                        {
                            m_ShaderInfos = shaderInfos;
                            m_OCResult = result;
                            Repaint();
                        }
                        EditorUtility.ClearProgressBar();
                    });
            }
        }
        
        m_ShowMode = GUILayout.Toolbar(m_ShowMode, new[] {new GUIContent("列表模式"), new GUIContent("文本模式")});
        if (m_ShowMode == 0)
        {
            ShowUIMode();
        }
        else if (m_ShowMode == 1)
        {
            ShowTextMode();
        }
    }
    
    private void ShowTextMode()
    {
        m_ScrollPos = EditorGUILayout.BeginScrollView(m_ScrollPos);

        EditorGUILayout.TextArea(m_OCResult);

        EditorGUILayout.EndScrollView();

        if (GUILayout.Button("保存"))
        {
            string path = EditorUtility.SaveFilePanel("保存", string.Empty, "Analysis", "txt");
            if (!string.IsNullOrEmpty(path))
            {
                File.WriteAllText(path, m_OCResult);
                AssetDatabase.Refresh();
            }
        }
    }
    
    private void ShowUIMode()
    {
        EditorGUILayout.BeginHorizontal();
        m_SearchPattern = EditorGUILayout.TextField("搜索Keyword", m_SearchPattern);
        if (GUILayout.Button("搜索", GUILayout.Width(44.0f)))
        {
            m_SearchPatterns.Clear();
            string[] patterns = m_SearchPattern.Split(new[] {' ', ','}, StringSplitOptions.RemoveEmptyEntries);
            m_SearchPatterns.AddRange(patterns);
        }

        EditorGUILayout.EndHorizontal();

        m_ScrollPos = EditorGUILayout.BeginScrollView(m_ScrollPos);

        int searchCount = 0;
        for (int i = 0; i < m_ShaderInfos.Count; ++i)
        {
            ShaderInfo shaderInfo = m_ShaderInfos[i];

            // 根据搜索框的内容，搜索Keyword
            if (m_SearchPatterns.Count > 0)
            {
                int foundCount = 0;
                for (int j = 0; j < m_SearchPatterns.Count; ++j)
                {
                    string p = m_SearchPatterns[j];
                    if (shaderInfo.m_LocalKeywords.Contains(p))
                    {
                        ++foundCount;
                    }
                    else if (shaderInfo.m_GlobalKeywords.Contains(p))
                    {
                        ++foundCount;
                    }
                }

                if (foundCount < m_SearchPatterns.Count)
                {
                    continue;
                }
            }

            ++searchCount;

            EditorGUILayout.BeginVertical("Button");
            shaderInfo.Expand = EditorGUILayout.Foldout(shaderInfo.Expand, "Shader变体：" + shaderInfo.m_ShaderName);
            if (!shaderInfo.Expand)
            {
                EditorGUILayout.EndVertical();
                continue;
            }

            EditorGUILayout.EndVertical();

            EditorGUILayout.BeginHorizontal();
            GUILayout.Space(22.0f);
            EditorGUILayout.BeginVertical();

            EditorGUILayout.BeginVertical("TextArea");

            EditorGUILayout.LabelField("类型", shaderInfo.m_ShaderType);

            #region 全局关键字

            EditorGUILayout.BeginHorizontal();
            EditorGUILayout.LabelField("全局关键字", GUILayout.Width(120.0f));
            EditorGUILayout.TextArea(shaderInfo.m_GlobalKeywordsText);
            EditorGUILayout.EndHorizontal();

            #endregion

            #region 局部关键字

            EditorGUILayout.BeginHorizontal();
            EditorGUILayout.LabelField("局部关键字", GUILayout.Width(120.0f));
            EditorGUILayout.TextArea(shaderInfo.m_LocalKeywordsText);
            EditorGUILayout.EndHorizontal();

            #endregion

            if (shaderInfo.m_DefaultInfo != null)
            {
                OnShaderBaseInfoGUI(shaderInfo.m_DefaultInfo, "默认变体");
            }

            if (shaderInfo.m_PositionVariant != null)
            {
                OnShaderBaseInfoGUI(shaderInfo.m_PositionVariant, "位置变体(Position Variant)");
            }

            if (shaderInfo.m_VaringVariant != null)
            {
                OnShaderBaseInfoGUI(shaderInfo.m_VaringVariant, "插值变体(Varing Variant)");
            }

            if (shaderInfo.m_ShaderProperties != null)
            {
                OnShaderPropertiesGUI(shaderInfo.m_ShaderProperties);
            }

            EditorGUILayout.EndVertical();

            EditorGUILayout.EndVertical();
            GUILayout.Space(22.0f);
            EditorGUILayout.EndHorizontal();

            EditorGUILayout.Space();
        }

        EditorGUILayout.EndScrollView();

        if (m_SearchPatterns.Count > 0)
        {
            EditorGUILayout.LabelField("搜索结果", string.Format("{0} 个变体", searchCount));
        }

        EditorGUILayout.BeginHorizontal();
        if (GUILayout.Button("折叠所有"))
        {
            for (int i = 0; i < m_ShaderInfos.Count; ++i)
            {
                ShaderInfo shaderInfo = m_ShaderInfos[i];
                shaderInfo.Expand = false;
            }
        }

        if (GUILayout.Button("展开所有"))
        {
            for (int i = 0; i < m_ShaderInfos.Count; ++i)
            {
                ShaderInfo shaderInfo = m_ShaderInfos[i];
                shaderInfo.Expand = true;
            }
        }

        EditorGUILayout.EndHorizontal();
    }
    
    private void OnShaderPropertiesGUI(ShaderProperties properties)
        {
            EditorGUILayout.LabelField("Shader属性：");
            EditorGUILayout.BeginHorizontal();
            GUILayout.Space(22.0f);
            EditorGUILayout.BeginVertical();
            {
                string tips;
                tips = "这是一种仅依赖字面量或者Uniform值的计算，因此在每一个线程都会产生相同的结果。虽然驱动程序可以优化这一点，但它仍然有成本，所以尽可能把它从着色器移到CPU逻辑中。uniform应该为常量。";
                EditorGUILayout.TextField(new GUIContent("是否有Uniform计算", tips), properties.m_HasUniformComputation.ToString());

                tips = "在固定图形管道之外，显示这个着色器是否有在内存中可见的副作用。它们可能由以下原因引起：\r\n" +
                       "● 写入着色器存储缓冲区（storage buffer）\r\n" +
                       "● 存储到图像中\r\n" +
                       "● 原子操作（atomics）的使用\r\n" +
                       "有副作用的着色器不能通过诸如隐藏表面消除（hidden surface removal, HSR）之类的技术来优化，所以它们的使用应该最小化。";
                EditorGUILayout.TextField(new GUIContent("是否有副作用", tips), properties.m_HasSideEffects.ToString());

                tips = "判断是否片段着色器包含覆盖掩码，可以通过着色器的执行来改变。例如discard操作。具有可修改覆盖严码的Shader，肯定会使用Late ZS Update，这降低了同一像素进行Early ZS Test的效率。";
                EditorGUILayout.TextField(new GUIContent("是否修改覆盖", tips), properties.m_ModifiesCoverage.ToString());

                tips = "片段着色器是否包含Late ZS Test操作，例如写了gl_FragDepth。Late ZS负责处理那些像素着色之前无法处理的深度、模板测试以及更新。";
                EditorGUILayout.TextField(new GUIContent("是否使用了Late ZS Test", tips), properties.m_UsesLateZSTest.ToString());

                tips = "片段着色器是否包含Late ZS Update逻辑，例如，通过使用gl_LastFragDepthARM读取着色器中的旧深度值。";
                EditorGUILayout.TextField(new GUIContent("是否使用了Late ZS Update", tips), properties.m_UseLateZSUpdate.ToString());

                tips = "片段着色器是否包含从颜色缓冲区（color buffer）读取的逻辑，例如，从gl_LastFragColorARM里读取。以这种方式读取的Shader，被视为透明的，不能用作隐藏表面消除（hidden surface removal, HSR）器。";
                EditorGUILayout.TextField(new GUIContent("是否读取颜色缓冲区", tips), properties.m_ReadColorBuffer.ToString());
            }
            EditorGUILayout.EndVertical();
            EditorGUILayout.EndHorizontal();
        }

        private void OnShaderBaseInfoGUI(ShaderBaseInfo baseInfo, string variantName)
        {
            EditorGUILayout.LabelField(variantName + "：");
            EditorGUILayout.BeginHorizontal();
            GUILayout.Space(22.0f);
            EditorGUILayout.BeginVertical();
            {
                EditorGUILayout.IntField("工作寄存器", baseInfo.m_WorkRegisters);
                EditorGUILayout.IntField("Uniform寄存器", baseInfo.m_UniformRegisters);
                EditorGUILayout.Toggle("是否有变量被放置到栈内存中", baseInfo.m_StackSpilling);
                EditorGUILayout.TextField("16位运算", $"{baseInfo.m_Arithmetic16Bit * 1e2f}%");

                EditorGUILayout.LabelField("Shader复杂度：");
                EditorGUILayout.BeginHorizontal();
                GUILayout.Space(22.0f);
                EditorGUILayout.BeginVertical();
                {
                    float valueWidth = 50.0f;

                    EditorGUILayout.BeginHorizontal();
                    EditorGUILayout.LabelField(" ", GUILayout.Width(160.0f));
                    EditorGUILayout.LabelField("A", GUILayout.Width(valueWidth));
                    EditorGUILayout.LabelField("LS", GUILayout.Width(valueWidth));
                    EditorGUILayout.LabelField("V", GUILayout.Width(valueWidth));
                    EditorGUILayout.LabelField("T", GUILayout.Width(valueWidth));
                    EditorGUILayout.LabelField("瓶颈", GUILayout.Width(valueWidth));
                    EditorGUILayout.EndHorizontal();

                    EditorGUILayout.BeginHorizontal();
                    EditorGUILayout.LabelField("指令的累积执行Cycles", GUILayout.Width(160.0f));
                    EditorGUILayout.FloatField(baseInfo.m_TotalShaderCycles.m_A, GUILayout.Width(valueWidth));
                    EditorGUILayout.FloatField(baseInfo.m_TotalShaderCycles.m_LS, GUILayout.Width(valueWidth));
                    EditorGUILayout.FloatField(baseInfo.m_TotalShaderCycles.m_V, GUILayout.Width(valueWidth));
                    EditorGUILayout.FloatField(baseInfo.m_TotalShaderCycles.m_T, GUILayout.Width(valueWidth));
                    EditorGUILayout.TextField(baseInfo.m_TotalShaderCycles.m_Bound, GUILayout.Width(valueWidth));
                    EditorGUILayout.EndHorizontal();

                    EditorGUILayout.BeginHorizontal();
                    EditorGUILayout.LabelField("最短控制流路径Cycles", GUILayout.Width(160.0f));
                    EditorGUILayout.FloatField(baseInfo.m_ShortestShaderCycles.m_A, GUILayout.Width(valueWidth));
                    EditorGUILayout.FloatField(baseInfo.m_ShortestShaderCycles.m_LS, GUILayout.Width(valueWidth));
                    EditorGUILayout.FloatField(baseInfo.m_ShortestShaderCycles.m_V, GUILayout.Width(valueWidth));
                    EditorGUILayout.FloatField(baseInfo.m_ShortestShaderCycles.m_T, GUILayout.Width(valueWidth));
                    EditorGUILayout.TextField(baseInfo.m_ShortestShaderCycles.m_Bound, GUILayout.Width(valueWidth));
                    EditorGUILayout.EndHorizontal();

                    EditorGUILayout.BeginHorizontal();
                    EditorGUILayout.LabelField("最长控制流路径Cycles", GUILayout.Width(160.0f));
                    EditorGUILayout.FloatField(baseInfo.m_LongestShaderCycles.m_A, GUILayout.Width(valueWidth));
                    EditorGUILayout.FloatField(baseInfo.m_LongestShaderCycles.m_LS, GUILayout.Width(valueWidth));
                    EditorGUILayout.FloatField(baseInfo.m_LongestShaderCycles.m_V, GUILayout.Width(valueWidth));
                    EditorGUILayout.FloatField(baseInfo.m_LongestShaderCycles.m_T, GUILayout.Width(valueWidth));
                    EditorGUILayout.TextField(baseInfo.m_LongestShaderCycles.m_Bound, GUILayout.Width(valueWidth));
                    EditorGUILayout.EndHorizontal();
                }
                EditorGUILayout.EndVertical();
                EditorGUILayout.EndHorizontal();
            }
            EditorGUILayout.EndVertical();
            EditorGUILayout.EndHorizontal();
        }
}